#! /bin/bash

RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
BLUE="\033[34m"
END="\033[0m"

# root权限
root_privilege() {
    if [[ $EUID -ne 0 ]]; then
        echo "Error: 请使用root权限运行该脚本!" 1>&2
        exit 1
    fi
}

# 系统信息
system_info() {
    clear
    ## 定义系统名称(Define system name)
    SYSTEM_NAME=$(cat /etc/os-release  | grep -E "^NAME=" | awk -F '=' '{print$2}' | sed "s/[\'\"]//g")
    ## 定义系统版本(Define system version)
    SYSTEM_VERSION_NUMBER=$(cat /etc/os-release  | grep -E "VERSION_ID=" | awk -F '=' '{print$2}' | sed "s/[\'\"]//g")
    ARCH=$(uname -m)
}

file_name=$1
param_count="$#"
# 判断参数个数并进行相应操作
if [ "$param_count" -ne 1 ]; then
    echo "未传递参数 或 存在多个参数，只需传递IP文件1个参数即可"
    exit 1
fi
ipset_name=$(echo $1 | awk -F "." '{print $1}' 2>/dev/null)

block_city() {
    check_ipset
    # 创建本地局域网
    local_ips=$(iptables -L INPUT --line-numbers | grep local_ips | wc -l)
    if [ $local_ips -eq 0 ]; then
        echo -ne "${YELLOW}\033[K 局域网IP，正在添加中...${END}\r"
        ipset create local_ips hash:net
        ipset add local_ips 10.0.0.0/8
        ipset add local_ips 172.16.0.0/12
        ipset add local_ips 192.168.0.0/16
        iptables -I INPUT -m set --match-set local_ips src -j ACCEPT
        echo -ne "${GREEN}\033[K 局域网IP，完成！${END}\r"
        echo ""
    fi

    sed -i 's/\r$//' *.txt
    # 添加规则
    cq=$(ipset list $ipset_name 2>/dev/null | wc -l)
    if [ $cq -eq 0 ]; then
        echo -ne "${YELLOW}\033[K $ipset_name，正在添加中...${END}\r"
        ipset create $ipset_name hash:net
        for ip_ranges in $(cat "$file_name"); do ipset add $ipset_name $ip_ranges ;done
    fi
    iptables_rules=$(iptables -L --line-numbers | grep match-set | grep $ipset_name | wc -l)
    if [ $iptables_rules -eq 0 ]; then
        iptables -I INPUT -m set --match-set $ipset_name src -j ACCEPT
        echo -e "${GREEN}\033[K $ipset_name，完成！${END}\r"
        echo ""
    fi
    cat > /usr/lib/systemd/system/ipset-restore-$ipset_name.service <<EOF
[Unit]
Description=ipset rules server
Wants=network-online.target
After=network.target

[Service]
Type=forking
WorkingDirectory=$(pwd)
ExecStart=/bin/bash my_ips.sh $file_name
Restart=on-failure

[Install]
WantedBy=multi-user.target
EOF
    systemctl enable ipset-restore-$ipset_name.service
}

# 删除iptables规则
unblock_ipset() {
    echo ""
    iptables -L INPUT --line-numbers | awk '/match-set/ {print $1, $8, $2}' | sed -E 's/(.*)/\x1b[34m\1\x1b[0m/'
    iptables -L INPUT --line-numbers | grep -P '^\d+\s+DROP\s+all\s+--\s+anywhere\s+anywhere\s*$' | sed -E 's/(.*)/\x1b[34m\1\x1b[0m/'
    iptables_list=$(iptables -L INPUT --line-numbers | grep 'match-set' | wc -l)
    if [ $iptables_list -eq 0 ]; then
        ipset destroy local_ips 2>/dev/null
        ipset destroy $ipset_name 2>/dev/null
        rm -rf /usr/lib/systemd/system/ipset-restore-$ipset_name.service
        exit 1
    fi
    echo -e "\n${BLUE}请输入规则列表中的需要删除的编号${END}"
    read -p "$(echo -e "\n${YELLOW}└─ Please enter the number： ${END}")" NUM_CODE
    
    if [[ $NUM_CODE =~ ^[0-9]+$ ]] && [[ $NUM_CODE -le $((iptables_list + 1)) ]]; then
        iptables -D INPUT $NUM_CODE
        ipset destroy $iptables_name
        echo -ne "${GREEN}\033[K $ipset_name 删除完成！${END}\r"
        echo ""
    else 
        echo -e "\n${RED} 输入有误，请重新输入！${END}\n"
        unblock_ipset
    fi
    rm -rf /usr/lib/systemd/system/ipset-restore-$ipset_name.service
}

# 查看规则列表
block_list() {
    echo ""
    iptables -L --line-numbers | grep match-set
    iptables -L INPUT --line-numbers | grep -P '^\d+\s+DROP\s+all\s+--\s+anywhere\s+anywhere\s*$'
    echo ""
}


# 检查ipset是否安装
check_ipset() {
    if [ -f /sbin/ipset ]; then
        echo -e ""
    elif [ $(echo SYSTEM_VERSION_NUMBER | grep Linux | wc -l) -ge 1 ]; then
        yum -y install ipset
    else
        apt-get -y install ipset
    fi
}


#开始菜单
main() {
root_privilege
system_info
echo -e ""
echo -e "+---------------------------------------------------+"
echo -e "|        \033[1m欢迎使用 IP 访问控制系统${END}                   |"
echo -e '+---------------------------------------------------+'
echo -e "         运行环境: ${BLUE}${SYSTEM_NAME} ${SYSTEM_VERSION_NUMBER}${END} "
echo -e "         系统时间: ${BLUE}$(date "+%Y-%m-%d %H:%M:%S")${END}            "
echo -e "+---------------------------------------------------+"
echo -e ""
echo -e "\033[36m         ❖  允许某城市(IP) \t  \t 1 )${END}"
echo -e "\033[36m         ❖  删除某城市(IP) \t  \t 2 )${END}"
echo -e "\033[36m         ❖  查看规则列表 \t \t 3 )${END}"

CHOICE_NUM=$(echo -e "\n${YELLOW}└─ 请输入数字编号? :  ${END}")
read -p "${CHOICE_NUM}" INPUT

case "$INPUT" in
    1)
    block_city
    ;;
    2)
    unblock_ipset
    ;;
    3)
    block_list
    ;;
    *)
    echo -e "${YELLOW}└─ 请输入数字编号： ${END}"
    main
    ;;
    esac
}
#main
block_city